#include <iostream>
#include <vector>
#include <ctime>

#include "augmentOpencv.h"
#include "KinectWrapper.h"
#include "Color2GrayMapper.h"
#include "DepthTracker.h"
#include "DataReader.h"

using namespace std;
using namespace cv;

const int WIDTH = 640;
const int HEIGHT = 480;

vector< vector<Point3f> > objPts;
vector< vector<Point2f> > imgPtsC;
vector< vector<Point2f> > imgPtsG;
int saveIdx = 0;

void saveCaliResult(string fn, vector< vector<Point3f> > objPts, Size size, vector< vector<Point2f> > imgPts, 
	vector<Mat> rvecs, vector<Mat> tvecs, Mat intriMat, Mat distortMat) 
{
	FileStorage fs(fn, FileStorage::WRITE);

	vector<float> flatVec;
	for (int i=0; i<objPts.size(); ++i) {
		for (int j=0; j<objPts[i].size(); ++j) {
			flatVec.push_back(objPts[i][j].x);
			flatVec.push_back(objPts[i][j].y);
			flatVec.push_back(objPts[i][j].z);
		}
		flatVec.push_back(-1);
		flatVec.push_back(-1);
		flatVec.push_back(-1);
	}
	Mat objPtsMat(flatVec.size()/3, 3, CV_32FC1, flatVec.data());
	fs << "objPtsMat" << objPtsMat;

	fs << "Size" << size;

	flatVec.clear();
	for (int i=0; i<imgPts.size(); ++i) {
		for (int j=0; j<imgPts[i].size(); ++j) {
			flatVec.push_back(imgPts[i][j].x);
			flatVec.push_back(imgPts[i][j].y);
			flatVec.push_back(1);
		}
		flatVec.push_back(-1);
		flatVec.push_back(-1);
		flatVec.push_back(-1);
	}
	Mat imgPtsMat(flatVec.size()/3, 3, CV_32FC1, flatVec.data());
	fs << "imgPtsMat" << imgPtsMat;

	fs << "rvecs" << "[";
	for (int i=0; i<rvecs.size(); ++i)
		fs << rvecs[i];
	fs << "]";

	fs << "tvecs" << "[";
	for (int i=0; i<tvecs.size(); ++i)
		fs << tvecs[i];
	fs << "]";

	fs << "intriMat" << intriMat;
	fs << "distortMat" << distortMat;
	fs.release();
}

void caliBoth(void) {
	KinectDepthColorWrapper kdcWrapper;
	CvMat *depthRes = NULL, *colorRes = NULL, *mappedRes = NULL;
	Mat depth, color, mapped, gray;

	kdcWrapper.init(WIDTH, HEIGHT);

	VideoCapture capture(CV_CAP_DSHOW+1);
	capture.set(CV_CAP_PROP_FRAME_WIDTH, WIDTH);
	capture.set(CV_CAP_PROP_FRAME_HEIGHT, HEIGHT);


	while (true) {
		kdcWrapper.retrieve(depthRes, colorRes, mappedRes);
		capture.read(gray);
		flip(gray, gray, 1);

		depth = Mat(depthRes);
		color = Mat(colorRes);
		mapped = Mat(mappedRes);

		imshow("depth", depth);
		imshow("color", color);
		imshow("mapped", mapped);
		imshow("gray", gray);

		char ch = cvWaitKey(50);
		if (ch == 27) break;

		if (ch == 13) {
			bool finish = false;

			for (int i=9; i>=5; --i) {
				if (finish) break;
				for (int j=i; j>=5; --j) {
					Size patternSize(i,j);
					vector<Point2f> cornersC;
					vector<Point2f> cornersG;
					if (!findChessboardCorners(color, patternSize, cornersC, 1+8)
						|| !findChessboardCorners(gray, patternSize, cornersG, 1+8)) {

							if (!findChessboardCorners(color, patternSize, cornersC)
								|| !findChessboardCorners(gray, patternSize, cornersG))
								continue;
					}

					cout << i << " " << j << endl;

					Mat drawImg = color;
					drawChessboardCorners(drawImg, patternSize, cornersC, 1);
					imshow("color_corner", drawImg);
					drawImg = gray;
					drawChessboardCorners(drawImg, patternSize, cornersG, 1);
					imshow("gray_corner", drawImg);

					ch = cvWaitKey(0);

					if (ch == 13) {
						cout << "save it..." << endl;
						finish = true;
						
						char prefix[100];
						sprintf(prefix, "data%d", saveIdx++);
						string prefixStr(prefix);
						
						imwrite(prefixStr+"_c.jpg", color);
						FileStorage fs(prefixStr+"_depthValue.yml", FileStorage::WRITE);
						fs << "mapped" << mapped;
						fs << "depth" << depth;
						fs.release();

						imwrite(prefixStr+"_g.jpg", gray);

						vector<Point3f> objpts;
						for (int a=0; a<j; ++a) {
							for (int b=0; b<i; ++b) {
								objpts.push_back(Point3f(b,a,0));
							}
						}
						objPts.push_back(objpts);
						imgPtsC.push_back(cornersC);
						imgPtsG.push_back(cornersG);

						break;
					}
				}
			}
		}
	}

	cout << "now calibration" << endl;
	Mat intriMatC, distortMatC, intriMatG, distortMatG;
	vector<Mat> rvecsC, tvecsC, rvecsG, tvecsG;
	calibrateCamera(objPts, imgPtsC, Size(HEIGHT, WIDTH), intriMatC, distortMatC, rvecsC, tvecsC);
	ezPrintMat(intriMatC);
	ezPrintMat(distortMatC);
	calibrateCamera(objPts, imgPtsG, Size(HEIGHT, WIDTH), intriMatG, distortMatG, rvecsG, tvecsG);
	ezPrintMat(intriMatG);
	ezPrintMat(distortMatG);

	Mat R,T,E,F;
	stereoCalibrate(objPts, imgPtsC, imgPtsG, intriMatC, distortMatC, intriMatG, distortMatG, 
		Size(HEIGHT,WIDTH), R, T, E, F);

	saveCaliResult("calibrate_color_result.yml", objPts, Size(HEIGHT,WIDTH), imgPtsC, rvecsC, tvecsC,
		intriMatC, distortMatC);
	saveCaliResult("calibrate_gray_result.yml", objPts, Size(HEIGHT,WIDTH), imgPtsG, rvecsG, tvecsG,
		intriMatG, distortMatG);
	
	FileStorage fs("stereo_result.yml", FileStorage::WRITE);
	fs << "R" << R;
	fs << "T" << T;
	fs << "E" << E;
	fs << "F" << F;
	
	fs.release();
}

void runMainLoop(void) {
	CvMat *depthRes = NULL, *colorRes = NULL, *mappedRes = NULL, *grayRes = NULL;
	Mat depth, color, mapped, gray;

#define USEFILE

#ifdef USEFILE
	DepthReader depthReader;
	ColorReader colorReader;
	GrayReader grayReader;
	grayRes = cvCreateMat(grayReader.height, grayReader.width, CV_8U);
	mappedRes = cvCreateMat(depthReader.height, depthReader.width, CV_16S);
	colorRes = cvCreateMat(colorReader.height, colorReader.width, CV_8UC3);
#else
	//init device
	KinectDepthColorWrapper kdcWrapper;
	kdcWrapper.init(WIDTH, HEIGHT);

	VideoCapture capture(CV_CAP_DSHOW+1);
	capture.set(CV_CAP_PROP_FRAME_WIDTH, WIDTH);
	capture.set(CV_CAP_PROP_FRAME_HEIGHT, HEIGHT);
#endif

	//init mapper
	Color2GrayMapper mapper;
	mapper.init();


	DepthTracker depthTracker;
	Rect targRect;
	Mat tiled;
	tiled = Mat(HEIGHT, 2*WIDTH, CV_8UC3);

	//video recorder
// 	VideoWriter vw;
// 	vw.open("tiled.avi", -1, 10, Size(2*WIDTH, HEIGHT));

	while (true) {

#ifdef USEFILE
		depthReader.getNextFrame(mappedRes->data.s);
		grayReader.getNextFrame(grayRes->data.ptr);
		colorReader.getNextFrame(colorRes->data.ptr);
		gray = Mat(grayRes);
#else
		kdcWrapper.retrieve(depthRes, colorRes, mappedRes);
		capture.read(gray);
		flip(gray, gray, 1);
#endif


		color = Mat(colorRes);
		mapped = Mat(mappedRes);
		

		targRect = depthTracker.tracking(mapped);

		Mat mapRes = mapper.mapping(color, mapped, targRect);

		rectangle(mapped, targRect, Scalar(32767));
		imshow("mapped", mapped);

		cout << "mapRes size: " << mapRes.rows << " " << mapRes.cols << endl;
		imshow("gray", gray);
 		imshow("map result", mapRes);
		//for recorder
// 		gray.copyTo(tiled.colRange(Range(0, gray.cols)));
// 		mapRes.copyTo(tiled.colRange(Range(gray.cols, gray.cols+mapRes.cols)));
// 		imshow("tiled", tiled);
// 		vw.write(tiled);

		char c = cvWaitKey(30);
		if (c == 27) break;
	}
	
}

void changeFPS(void) {
	VideoCapture vc("tiled.avi");
	VideoWriter vw;
	vw.open("tiled_.avi", -1, 10, Size(2*WIDTH, HEIGHT));
	Mat img;
	while (vc.read(img)) {
		imshow("img", img);
		vw.write(img);
		cvWaitKey(30);
	}
}

void findRotation(void) {
	KinectDepthColorWrapper kdcWrapper;
	CvMat *depthRes = NULL, *colorRes = NULL, *mappedRes = NULL;
	Mat depth, color, mapped, gray;

	kdcWrapper.init(WIDTH, HEIGHT);

	VideoCapture capture(CV_CAP_DSHOW+1);
	capture.set(CV_CAP_PROP_FRAME_WIDTH, WIDTH);
	capture.set(CV_CAP_PROP_FRAME_HEIGHT, HEIGHT);

	Mat objPts, imgPtsC, imgPtsG;

	while (true) {
		kdcWrapper.retrieve(depthRes, colorRes, mappedRes);
		capture.read(gray);
		flip(gray, gray, 1);

		depth = Mat(depthRes);
		color = Mat(colorRes);
		mapped = Mat(mappedRes);

		imshow("depth", depth);
		imshow("color", color);
		imshow("mapped", mapped);
		imshow("gray", gray);

		char ch = cvWaitKey(0);
		if (ch == 27) break;

		if (ch == 13) {
			bool finish = false;

			for (int i=6; i>=3; --i) {
				if (finish) break;
				for (int j=i; j>=3; --j) {
					Size patternSize(i,j);
					vector<Point2f> cornersC;
					vector<Point2f> cornersG;
					if (!findChessboardCorners(color, patternSize, cornersC, 1+8)
						|| !findChessboardCorners(gray, patternSize, cornersG, 1+8)) {

							if (!findChessboardCorners(color, patternSize, cornersC)
								|| !findChessboardCorners(gray, patternSize, cornersG))
								continue;
					}

					cout << i << " " << j << endl;

					Mat drawImg = color;
					drawChessboardCorners(drawImg, patternSize, cornersC, 1);
					imshow("color_corner", drawImg);
					drawImg = gray;
					drawChessboardCorners(drawImg, patternSize, cornersG, 1);
					imshow("gray_corner", drawImg);

					ch = cvWaitKey(0);

					if (ch == 13) {
						finish = true;

						vector<Point3f> objpts;
						for (int a=0; a<j; ++a) {
							for (int b=0; b<i; ++b) {
								objpts.push_back(Point3f(b,a,0));
							}
						}

						objPts = Mat(objpts.size(), 3, CV_64F);
						for (int i=0; i<objpts.size(); ++i) {
							objPts.at<double>(i,0) = objpts[i].x;
							objPts.at<double>(i,1) = objpts[i].y;
							objPts.at<double>(i,2) = objpts[i].z;
						}

						imgPtsC = Mat(cornersC.size(), 2, CV_64F);
						for (int i=0; i<cornersC.size(); ++i) {
							imgPtsC.at<double>(i,0) = cornersC[i].x;
							imgPtsC.at<double>(i,1) = cornersC[i].y;
						}

						imgPtsG = Mat(cornersG.size(), 2, CV_64F);
						for (int i=0; i<cornersG.size(); ++i) {
							imgPtsG.at<double>(i,0) = cornersG[i].x;
							imgPtsG.at<double>(i,1) = cornersG[i].y;
						}

						break;
					}
				}
			}
		
			if (finish) break;
			cout << "fail" << endl;
		}
	}

	Mat intriMatC, distortMatC, intriMatG, distortMatG;

	FileStorage fs;
	fs.open("calibrate_color_result.yml", FileStorage::READ);
	fs["intriMat"] >> intriMatC;
	fs["distortMat"] >> distortMatC;
	fs.release();
	fs.open("calibrate_gray_result.yml", FileStorage::READ);
	fs["intriMat"] >> intriMatG;
	fs["distortMat"] >> distortMatG;
	fs.release();

	CvMat cvintriMatC = intriMatC, cvdistortMatC = distortMatC;
	CvMat cvintriMatG = intriMatG, cvdistortMatG = distortMatG;
	CvMat cvobjPts = objPts, cvimgPtsG = imgPtsG, cvimgPtsC = imgPtsC;
	CvMat *rvecC = cvCreateMat(3,1,CV_64F);
	CvMat *tvecC = cvCloneMat(rvecC);
	CvMat *rvecG = cvCloneMat(rvecC);
	CvMat *tvecG = cvCloneMat(tvecC);

 	cvFindExtrinsicCameraParams2(&cvobjPts, &cvimgPtsC, &cvintriMatC, &cvdistortMatC,
		rvecC, tvecC);
	cvFindExtrinsicCameraParams2(&cvobjPts, &cvimgPtsG, &cvintriMatG, &cvdistortMatG,
		rvecG, tvecG);

	Mat RG, RC, TG, TC;
	Rodrigues(Mat(rvecG), RG);
	Rodrigues(Mat(rvecC), RC);
	TG = Mat(tvecG);
	TC = Mat(tvecC);

	Mat R,T,E,F;
	fs.open("stereo_result.yml", FileStorage::READ);
	fs["R"] >> R;
	fs["T"] >> T;
	fs["E"] >> E;
	fs["F"] >> F;
	fs.release();

	R = RG*RC.inv();
	T = -R*TC+TG;

	fs.open("stereo_result.yml", FileStorage::WRITE);
	fs << "R" << R;
	fs << "T" << T;
	fs << "E" << E;
	fs << "F" << F;
	fs << "edgeLen" << 15.5;
	fs.release();
}

int main(void)
{
	char mode = 'r';
	cout << "choose system working mode: (c)calibrate or (r)run or (f)find rotation" << endl;
	cin >> mode;
	cout << "chose mode " << mode << endl;
// 	changeFPS();
	if (mode == 'c') {
		caliBoth();
	}
	else if (mode == 'r') {
		runMainLoop();
	}
	else if (mode == 'f') {
		findRotation();
	}

	return 0;
}